#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
密码管理器登录的功能。
'''
__author__ = 'chenwei'
__version__ = 'v1.0.0'
import os
import pickle
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex
import random
passwd_key = "password_manager" #prpcrypt里面需要，这个变量必须为16位。否则会报错
file_dir = "file" #文件目录
main_passwd_db = os.path.join(file_dir, "main_passwd.db") #文件存放的相对位置。
web_passwd_db = os.path.join(file_dir, "web_passwd.db")
bank_passwd_db = os.path.join(file_dir, "bank_passwd.db")
app_passwd_db = os.path.join(file_dir, "app_passwd.db")
server_passwd_db = os.path.join(file_dir, "server_passwd.db")


class prpcrypt():
    def __init__(self, key):
        self.key = key
        self.mode = AES.MODE_CBC

    # 加密函数，如果text不是16的倍数【加密文本text必须为16的倍数！】，那就补足为16的倍数
    def encrypt(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        # 这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.目前AES-128足够用
        length = 16
        count = len(text)
        add = length - (count % length)
        text = text + ('\0' * add)
        self.ciphertext = cryptor.encrypt(text)
        # 因为AES加密时候得到的字符串不一定是ascii字符集的，输出到终端或者保存时候可能存在问题
        # 所以这里统一把加密后的字符串转化为16进制字符串
        return str(b2a_hex(self.ciphertext),encoding="utf-8")

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        plain_text = cryptor.decrypt(a2b_hex(text))
        return str(plain_text,encoding = "utf-8").rstrip('\0')

    def encrypt_hex(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        # 这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.目前AES-128足够用
        text2 = str(b2a_hex(text), encoding="utf-8")
        length = 16
        count = len(text2)
        add = length - (count % length)
        text2 = text2 + ('\0' * add)
        self.ciphertext = cryptor.encrypt(text2)
        return self.ciphertext

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt_hex(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        plain_text = cryptor.decrypt(text)
        return str(plain_text,encoding = "utf-8").rstrip('\0')

def open_file(fn, md='rb'):  # 打开文件函数
    try:
        return open(fn, md)
    except FileNotFoundError as e:
        return open(fn, 'w')
    except IOError as e:
        raise e


def load_pickle(fn):#读取序列化文件。
    with open_file(fn) as f:
        try:
            enc_pkl_item = f.read()
            pc = prpcrypt(passwd_key)  # 初始化密钥
            pkl_item = pc.decrypt_hex(enc_pkl_item)
            pkl_item = a2b_hex(pkl_item)
            return pickle.loads(pkl_item)
        except EOFError as e:
            raise e
        except Exception as e:
            raise e


def dump_pickle(item,fn):#保存序列化文件。
    with open_file(fn, md='w+b') as f:
        try:
            pkl_item = pickle.dumps(item)
            pc = prpcrypt(passwd_key)  # 初始化密钥
            enc_pkl_item = pc.encrypt_hex(pkl_item)
            f.write(enc_pkl_item)
        except Exception as e:
            raise e
    return True


def input_password():
    passwd1 = input("请输入密码:")
    passwd2 = input("请再次输入密码:")
    if passwd1 != passwd2:
        return False
    return passwd1

def input_authority():#权限修改，是管理员则返回True，不是管理员则返回Flase。
    authority = input("是否为管理用户？(y/n)默认为n")
    if authority != 'y' or authority != 'Y':
        return True
    return False


def print_menu(this_list): #打印菜单菜单
    print(this_list,end='')


class login_acction:#为多用户设计的登录界面
    def __init__(self,fn):
        self.__data = {}
    @property
    def data(self):
        return self.__data
    def modify_password(self,username):
        result = input_password()
        if not result:
            return False
        self.__data[username] = {'passwd':result}
        return True
    def modify_authority(self,username):
        result = input_authority()
        if not result:
            self.__data[username]['auth'] = False
        self.__data[username]['auth'] = True
        return True
    def save_file(self,fn):
        dump_pickle(self,fn)
        return True
    def check_passwd(self,username,passwd):#验证密码是否正确，正确返回True，错误返回False
        result = self.check_username(username)
        if not result:
            return False
        if self.__data[username]['passwd'] != passwd:
            return False
        if self.__data[username]['passwd'] == passwd:
            return True
    def login(self,username,passwd):#登录模块，以后会实现验证权限，验证是否为删除用户。
        return self.check_passwd(username,passwd)
    def check_username(self,username):
        if username not in self.__data:# 用户中是否有同名用户,有则返回真，没有返回假
            return False
        return True
    def creater_user(self,username):
        result = self.check_username(username)
        print(result)
        if result:#如果查到同名用户，则返回False
            return False
        result = self.modify_password(username)
        print(result)
        if not result:#如果两次输入的密码不相同，则返回假。
            return False
        self.modify_authority(username)#修改用户的权限，
        print(result)
        if not result:
            return False
        return True

def myAlign(string, length=0,placeholder= ' ',head=False,head_placeholder= ' ',tail=False,tail_placeholder= ' '):
    if length == 0:
        return string
    slen= len(string)
    re = string
    flag = False
    while slen < length-2:
        flag = True
        if slen % 2 == 0:
            re += placeholder
        else:
            re = placeholder + re
        slen += 1
    if head:
        re = head_placeholder + re
    elif flag:
        re = placeholder + re
    if tail:
        re += tail_placeholder
    elif flag:
        re += placeholder
    return re

def generate_verification_code(len=10):
    ''' 随机生成X位的验证码 '''
    # 注意： 这里我们生成的是0-9A-Za-z的列表，当然你也可以指定这个list，这里很灵活
    # 比如： code_list = ['P','y','t','h','o','n','T','a','b'] # PythonTab的字母
    code_list = []
    for i in range(10): # 0-9数字
        code_list.append(str(i))
    for i in range(65, 91): # 对应从“A”到“Z”的ASCII码
        code_list.append(chr(i))
    for i in range(97, 123): #对应从“a”到“z”的ASCII码
        code_list.append(chr(i))
    code_list.append('@')
    code_list.append('#')
    code_list.append('$')
    myslice = random.sample(code_list, len)  # 从list中随机获取6个元素，作为一个片断返回
    verification_code = ''.join(myslice) # list to string
    return verification_code


class passwd_manager(dict):
    def __init__(self):
        super().__init__()
        self.data = {}
        self.catagory = ''
    def insert_data(self, title, username, password, url):
        self.data[title] = [username, password, url]
    def modify_data(self, title, username, password, url):
        self.data[title] = [username, password, url]
    def del_data(self,title):
        try:
            self.data.pop(title)
        except KeyError as e:
            return False
    def find_data(self,title):
        tmp={}
        for key in self.data.keys():
            if title in key:
                tmp[key]=self.data[key]
        if tmp == {}:
            print('Not Found!')
            return False
        self.print_dict(obj=tmp.items())
    def find_tittle(self,title):
        return self.data.get(title,False)
    def print_dict(self,obj='',col_1='title', col_2='username', col_3='password', col_4='url'):
        if obj == '':
            obj = self.data.items()
        print(myAlign(self.catagory,100))#表头的输出
        print(myAlign(col_1,20,'-',head=True,head_placeholder='+')\
              + myAlign(col_2,20,'-',head=True,head_placeholder='+')\
              + myAlign(col_3,20,'-',head=True,head_placeholder='+')\
              + myAlign(col_4,40,'-',head=True,head_placeholder='+',tail=True,tail_placeholder='+'))#格式化输出报表头
        for key,value in obj:
            print(myAlign(key,20,' ',head=True,head_placeholder='|') \
                  + myAlign(value[0], 20, ' ', head=True, head_placeholder='|') \
                  + myAlign(value[1], 20, ' ', head=True, head_placeholder='|') \
                  + myAlign(value[2], 40, ' ', head=True, head_placeholder='|',tail=True,tail_placeholder='|'))#循环输出
        print(myAlign('',100,'-',head=True,head_placeholder='+',tail=True,tail_placeholder='+'))


class web_passwd_manager(passwd_manager):
    def __init__(self):
        super().__init__()
        self.data = {}
        self.catagory = 'web'
    def print_all(self):
        self.print_dict()
    def save_file(self):
        dump_pickle(self,web_passwd_db)
    @classmethod
    def get_item(self):
        return ['title', 'user', 'passwd', 'url']

class bank_passwd_manager(passwd_manager):
    def __init__(self):
        super().__init__()
        self.data = {}
        self.catagory = 'bank'
    def print_all(self):
        self.print_dict()
    def save_file(self):
        dump_pickle(self,bank_passwd_db)
    @classmethod
    def get_item(self):
        return ['title', 'user', 'passwd', 'url']

class server_passwd_manager(passwd_manager):
    def __init__(self):
        super().__init__()
        self.data = {}
        self.catagory = 'server'
    def print_all(self,col_4='ip:port'):
        self.print_dict(col_4=col_4)
    def save_file(self):
        dump_pickle(self,server_passwd_db)
    @classmethod
    def get_item(self):
        return ['title', 'user', 'passwd', 'ip:port']


class app_passwd_manager(passwd_manager):
    def __init__(self):
        super().__init__()
        self.data = {}
        self.catagory = 'app'
    def print_all(self):
        self.print_dict()
    def save_file(self):
        dump_pickle(self,app_passwd_db)
    @classmethod
    def get_item(self):
        return ['title', 'user', 'passwd', 'url']

def main():
    user_login = load_pickle(main_passwd_db)
    count = 0
    menu_list='''欢迎使用，密码管理器。'''

    while count < 3:
        print_menu("欢迎您使用，密码管理器。\n")
        username = input("请输入用户名:")
        password = input("请输入密码:")
        result = user_login.login(username,password)
        if not result:
            count += 1
            print("用户名或密码错误，请重新输入。最多尝试三次")
        else:
            cpass()
    else:
        print('错误次数过多，退出！')

def init_mainpasswd():
    a = login_acction(main_passwd_db)
    a.creater_user('admin')
    a.save_file(main_passwd_db)

def cpass():
    method = ''
    web = load_pickle(web_passwd_db) if not load_pickle(web_passwd_db) else web_passwd_manager()
    bank = load_pickle(bank_passwd_db) if not load_pickle(bank_passwd_db) else bank_passwd_manager()
    server = load_pickle(server_passwd_db) if not load_pickle(server_passwd_db) else server_passwd_manager()
    app = load_pickle(app_passwd_db) if not load_pickle(app_passwd_db) else app_passwd_manager()
    menu_list ='''\n----用户菜单----
|1)创建新密码         |
|2)查询密码           |
|3)修改密码           |
|4)删除密码           |
|5)保存退出           |
|6)不保存退出         |
----------------
请输入数字(1-4)或(q)：'''
    while True:
        print_menu(menu_list)
        main_method = input()
        if main_method == '1':
            while True:
                method = input('请输入要创建的分类:1）web 2）bank 3)server 4)app 5)返回上级菜单：')
                if method == '5':
                    break
                args= []
                random_method = ''
                random_method = input('是否使用随机密码?(y/n)')
                if random_method == 'y':
                    # random_lens = input('请输入随机密码长度:')
                    random_passwd = generate_verification_code()
                    print('生成的随机密码为:'+ random_passwd)
                if method == '1':
                    for i in web_passwd_manager.get_item():
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    web.insert_data(args[0],args[1],args[2],args[3])
                if method == '2':
                    for i in bank_passwd_manager.get_item():
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    bank.insert_data(args[0],args[1],args[2],args[3])
                if method == '3':
                    for i in server_passwd_manager.get_item():
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    server.insert_data(args[0],args[1],args[2],args[3])
                if method == '4':
                    for i in app_passwd_manager.get_item():
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    app.insert_data(args[0],args[1],args[2],args[3])
            else:
                print('输入有误！，请重新输入！')
        if main_method == '2':
            while True:
                method = input('请输入要查询的分类:1）web 2）bank 3)server 4)app 5）全部 6)退出：')
                if method == '6':
                    break
                random_method = ''
                keyword_method = input('是否通过关键字进行模糊查找？(y/n)')
                if keyword_method == 'y':
                    keyword= input('请输入关键字:')
                if method == '1':
                    if keyword_method ==  'y':
                        web.find_data(keyword)
                    else:
                        web.print_all()
                if method == '2':
                    if keyword_method ==  'y':
                        bank.find_data(keyword)
                    else:
                        bank.print_all()
                if method == '3':
                    if keyword_method ==  'y':
                        server.find_data(keyword)
                    else:
                        server.print_all()
                if method == '4':
                    if keyword_method ==  'y':
                        app.find_data(keyword)
                    else:
                        app.print_all()
                if method == '5':
                    if keyword_method ==  'y':
                        web.find_data(keyword)
                        bank.find_data(keyword)
                        server.find_data(keyword)
                        app.find_data(keyword)
                    else:
                        web.print_all()
                        bank.print_all()
                        server.print_all()
                        app.print_all()
            else:
                print('输入有误！，请重新输入！')
        if main_method == '3':
            while True:
                method = input('请输入要修改的分类:1）web 2）bank 3)server 4)app 5)返回上级菜单：')
                if method == '5':
                    break
                title_method = input('请输入要修改的title名：')
                args= []
                random_method = ''
                random_method = input('是否使用随机密码替换原有密码?(y/n)')
                if random_method == 'y':
                    # random_lens = input('请输入随机密码长度:')
                    random_passwd = generate_verification_code()
                    print('生成的随机密码为:'+ random_passwd)
                if method == '1':
                    result=web.find_tittle(title_method)
                    if result == False:
                        print('未找到对应的title！')
                        continue
                    for i in web_passwd_manager.get_item():
                        if i == 'title':
                            args.append(title_method)
                            continue
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    web.modify_data(args[0],args[1],args[2],args[3])
                if method == '2':
                    result = bank.find_tittle(title_method)
                    if result == False:
                        print('未找到对应的title！')
                        continue
                    for i in bank_passwd_manager.get_item():
                        if i == 'title':
                            args.append(title_method)
                            continue
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    bank.modify_data(args[0],args[1],args[2],args[3])
                if method == '3':
                    result = server.find_tittle(title_method)
                    if result == False:
                        print('未找到对应的title！')
                        continue
                    for i in server_passwd_manager.get_item():
                        if i == 'title':
                            args.append(title_method)
                            continue
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    server.modify_data(args[0],args[1],args[2],args[3])
                if method == '4':
                    result = app.find_tittle(title_method)
                    if result == False:
                        print('未找到对应的title！')
                        continue
                    for i in app_passwd_manager.get_item():
                        if i == 'title':
                            args.append(title_method)
                            continue
                        if i == 'passwd' and random_method == 'y':
                            args.append(random_passwd)
                            continue
                        args.append(input(i + ":"))
                    app.modify_data(args[0],args[1],args[2],args[3])

            else:
                print('输入有误！，请重新输入！')
        if main_method == '4':
            while True:
                method = input('请输入要删除的分类:1）web 2）bank 3)server 4)app 5)返回上级菜单：')
                if method == '5':
                    break
                title_method = input('请输入要删除的title名：')
                if method == '1':
                    web.del_data(title_method)
                if method == '2':
                    bank.del_data(title_method)
                if method == '3':
                    server.del_data(title_method)
                if method == '4':
                    app.del_data(title_method)
            else:
                print('输入有误！，请重新输入！')
        if main_method == '5':
            web.save_file()
            bank.save_file()
            server.save_file()
            app.save_file()
            exit()
        if main_method == '6':
            print('再见！')
            exit()
if __name__ ==  '__main__' : #程序入口
    main()



